#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据获取模块
用于从网络获取A股上证指数的历史数据
"""

import akshare as ak
import pandas as pd


def get_sh_index_data(period="daily", start_date="20200101", end_date="20231231"):
    """
    获取上证指数历史数据
    
    参数:
    period: 数据周期，daily表示日线数据
    start_date: 开始日期，格式为YYYYMMDD
    end_date: 结束日期，格式为YYYYMMDD
    
    返回:
    DataFrame: 包含上证指数历史数据的DataFrame
    """
    # 使用akshare库获取上证指数历史数据
    # 使用stock_zh_index_daily_em函数获取东方财富的指数历史数据
    try:
        stock_zh_index_daily_df = ak.stock_zh_index_daily_em(symbol="sh000001")
        
        # 筛选指定日期范围内的数据
        # 注意：新接口返回的数据列名可能不同，需要做相应调整
        stock_zh_index_daily_df['date'] = pd.to_datetime(stock_zh_index_daily_df['date'])
        
        start_date_dt = pd.to_datetime(start_date)
        end_date_dt = pd.to_datetime(end_date)
        
        mask = (stock_zh_index_daily_df['date'] >= start_date_dt) & (stock_zh_index_daily_df['date'] <= end_date_dt)
        filtered_data = stock_zh_index_daily_df.loc[mask]
        
        return filtered_data
    except Exception as e:
        print(f"获取数据时出错: {e}")
        # 如果上述方法失败，尝试另一种方法
        try:
            # 使用全量历史数据接口
            stock_zh_index_hist_df = ak.index_zh_a_hist(symbol="000001", period="daily", start_date=start_date, end_date=end_date)
            # 重命名列以匹配原有代码
            stock_zh_index_hist_df = stock_zh_index_hist_df.rename(columns={
                '日期': 'date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '成交额': 'amount'
            })
            return stock_zh_index_hist_df[['date', 'open', 'close', 'high', 'low', 'volume', 'amount']]
        except Exception as e2:
            print(f"尝试备用方法时也出错: {e2}")
            raise e


def calculate_moving_average(data, window=5):
    """
    计算移动平均线
    
    参数:
    data: 股票数据DataFrame
    window: 计算移动平均的窗口大小
    
    返回:
    Series: 移动平均值序列
    """
    # 使用pandas的rolling函数计算移动平均
    ma = data['close'].rolling(window=window).mean()
    return ma